/*
 *  Copyright (C) 2005 Alexandre Norman <norman@xael.org>
 *  Copyright (C) 2007 Gianluigi Tiesi <sherpya@netfarm.it>
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

/* Tested with python 2.5 and clamav 0.90.2
 * Should work with any version of python from 2.1
 */

#include <Python.h>
#include <clamav.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <fcntl.h>

#define PYCLAMAV_VERSION "0.4.1-ClamWin"

/* ********************************************************* */

/* For python prior to 2.3 */
#ifndef PyMODINIT_FUNC
#define PyMODINIT_FUNC void
#endif

#ifndef S_ISLNK
#define S_ISLNK(m) (0)
#endif

/* ********************************************************* */

/* Exception object for python */
static PyObject *PyclamavError;

/* Signature number */
unsigned int signumber = 0;

/* Structures for clamav */
struct cl_node *root = NULL;
struct cl_limits limits;

int vmain = 0, vdaily = 0;

/* Custom Db path */
char dbpath[1024];

void get_db_version(int *main_version, int *daily_version)
{
    const char *dbdir;
    char *path = NULL;
    struct cl_cvd *cvd;
    size_t len = 0;

    /* Database directory */
    dbdir = dbpath[0] ? dbpath : cl_retdbdir();

    /* Main */
    len = strlen(dbdir) + 22;
    path = (char *) malloc(len);
    snprintf(path, len, "%s/main.cvd", dbdir);
    if(access(path, 0) < 0)
        snprintf(path, len, "%s/main.inc/main.info", dbdir);
    if((cvd = cl_cvdhead(path)))
    {
        *main_version = cvd->version;
        cl_cvdfree(cvd);
    }
    free(path);

    /* Daily */
    len = strlen(dbdir) + 22;
    path = (char *) malloc(len);
    snprintf(path, len, "%s/daily.cvd", dbdir);
    if(access(path, 0) < 0)
        snprintf(path, len, "%s/daily.inc/daily.info", dbdir);
    if((cvd = cl_cvdhead(path)))
    {
        *daily_version = cvd->version;
        cl_cvdfree(cvd);
    }
    free(path);
}

int loaddb(void)
{
    int ret = 0;

    vmain = 0;
    vdaily = 0;
    /* so loaddb can reused */
    if (root)
    {
        cl_free(root);
        root = NULL;
    }

    if((ret = cl_load(dbpath[0] ? dbpath : cl_retdbdir(), &root, &signumber, CL_DB_STDOPT)))
        return ret;

    /* build the final tree */
    if((ret = cl_build(root)))
    {
        /* free the partial tree */
        cl_free(root);
        return ret;
    }

    get_db_version(&vmain, &vdaily);
    return 0;
}

/*
 * If the virus database has been changed, then
 * free the current tree and reload the new one
 */
int if_database_have_changed_then_reload(void)
{
    int main_version = 0;
    int daily_version = 0;

    /* Test if the database have changed or not yet loaded */
    /* If yes : reload DB */
    get_db_version(&main_version, &daily_version);

    if (!root || (main_version != vmain) || (daily_version != vdaily))
        return loaddb();
    return 0;
}

/*
 * Return pyclamav version
 */
static PyObject *pyclamav_version(PyObject *self, PyObject *args)
{
    return Py_BuildValue("s", PYCLAMAV_VERSION);
}

/*
 * Return clamav version
 */
static PyObject *pyclamav_get_version(PyObject *self, PyObject *args)
{
    const char *version;
    /* Clamav version */
    version = cl_retver();
    return Py_BuildValue("(s,i,i)", version, vmain, vdaily);
}

/*
 * Return number of signature in the database as integer
 */
static PyObject *pyclamav_get_numsig(PyObject *self, PyObject *args)
{
    /* Raise exception if database error */
    if (if_database_have_changed_then_reload())
        return NULL;
    return Py_BuildValue("i", signumber);
}

/*
 * Set libclamav debug mode
 */
static PyObject *pyclamav_set_debug(PyObject *self, PyObject *args)
{
    cl_debug();
    Py_INCREF(Py_None);
    return Py_None;
}

/*
 * Scan a file given as a parameter
 */
static PyObject *pyclamav_scanfile(PyObject *self, PyObject *args)
{
    char *file_to_scan;
    struct stat buf;
    unsigned long int size = 0;
    const char *virname;
    int ret = 0;

    if ((ret = if_database_have_changed_then_reload()))
    {
        PyErr_SetString(PyExc_ValueError,  cl_strerror(ret));
        return NULL;
    }

    if (!PyArg_ParseTuple(args, "s", &file_to_scan))
    {
        PyErr_SetString(PyExc_TypeError, "Pass filename to scan (string) as argument");
        return NULL;
    }

    if (stat(file_to_scan, &buf) < 0)
    {
        PyErr_SetString(PyExc_ValueError, strerror(errno));
        return NULL;
    }

    if (!(S_ISREG(buf.st_mode) || S_ISLNK(buf.st_mode)))
    {
        PyErr_SetString(PyExc_TypeError, "Not a regular file");
        return NULL;
    }

    ret = cl_scanfile(file_to_scan, &virname, &size, root, &limits, CL_SCAN_STDOPT);

    /* Test return code */
    switch (ret)
    {
        case CL_VIRUS : /* File contains a virus */
            return Py_BuildValue("(i,s)", 1, virname);
        case CL_CLEAN : /* File is clean */
            return Py_BuildValue("(i,s)", 0, "");
        default: /* Error : raise exception with message */
            PyErr_SetString(PyExc_ValueError,  cl_strerror(ret));
        return NULL;
    }
}

static PyObject *pyclamav_setdbpath(PyObject *self, PyObject *args)
{
    char *newdbpath = NULL;
    if (!PyArg_ParseTuple(args, "s", &newdbpath))
    {
        /* Raise exception with error message */
        PyErr_SetString(PyclamavError, "DB path must be a String");
        return NULL;
    }
    dbpath[0] = 0;
    strncpy(dbpath, newdbpath, sizeof(dbpath));
    Py_INCREF(Py_None);
    return Py_None;
}

static PyObject *pyclamav_loaddb(PyObject *self, PyObject *args)
{
    int ret = 0;
    PyObject *result = NULL;

    if (!PyArg_UnpackTuple(args, "loaddb", 0, 1, &result))
        return NULL;

    /* Optional alternative db path */
    if (result)
    {
        if (PyString_Check(result))
        {
            dbpath[0] = 0;
            strncpy(dbpath, PyString_AsString(result), sizeof(dbpath));
        }
        else
        {
            PyErr_SetString(PyclamavError, "DB path must be a String");
            return NULL;
        }
    }

    if ((ret = loaddb()))
    {
        PyErr_SetString(PyclamavError, cl_strerror(ret));
        return NULL;
    }

    Py_INCREF(Py_None);
    return Py_None;
}

static PyMethodDef ClamavMethods[] = {
    { "loaddb",      pyclamav_loaddb, METH_VARARGS|METH_KEYWORDS, "loaddb(path): Load db and init libclamav enigne.\n"
                      "If path is specified also set db path\n" },
    { "setdbpath",   pyclamav_setdbpath, METH_VARARGS, "setdbpath(path): Set database path.\n" },
    { "scanfile",    pyclamav_scanfile, METH_VARARGS, "scanfile(filename) : Scan a file for virus.\n"
                      "Arguments : filename (string)\n"
                      "Return a tupple (status, virusname) where status=0 when no virus found\n"
                      "or status=1 if a virus was found\n"
                      "May raise a ValueError exception if an error occurs\n"
                      "May raise a TypeError exception if wrong arguments are passed\n" },
    { "get_numsig",  pyclamav_get_numsig, METH_VARARGS, "get_numsig() : Get the number of known virus signatures\n"
                      "Arguments : None\n Return the number of known signatures.\n" },
    { "get_version", pyclamav_get_version, METH_VARARGS, "get_version() : Get Clamav version.\n"
                      "Arguments : None\n Return the version of Clamav as a tupple (version, main_version, daily_version).\n" },
    { "set_debug",   pyclamav_set_debug, METH_VARARGS, "set_debug(): Enable debug messages in libclamav.\nArguments : None\nReturn None.\n" },
    { "version",     pyclamav_version, METH_VARARGS, "version() : Get pyclamav version.\nArguments : None\n Return the version of pyclamav.\n" },
    { NULL, NULL, 0, NULL }
};

PyMODINIT_FUNC initpyclamav(void)
{
    PyObject *module, *dict;
    module = Py_InitModule("pyclamav", ClamavMethods);
    dict = PyModule_GetDict(module);

    PyclamavError = PyErr_NewException("pyclamav.error", NULL, NULL);
    PyDict_SetItemString(dict, "error", PyclamavError);

    /* Set documentation string for the module */
    PyDict_SetItemString(dict, "__doc__",
        PyString_FromString("pyclamav :\n\n This is a python binding to the C libclamav library\n"
                            "(from the Clamav project - http://www.clamav.net).\n"
                            "It can be used to easily allow a Python script to scan\n"
                            "a file or a buffer against known viruses.\n\n"
                            "Author : Alexandre Norman [norman@xael.org]\n"
                            "Author : Gianluigi Tiesi [sherpya@netfarm.it]\n\n"
                            "Functions :\n"
                            " - loaddb(string path) : Load a virus db and init db engine\n"
                            " - setdbpath(string path) : Set Virus Database path\n"
                            " - scanfile(string filename) : Scan a file for virus.\n"
                            " - get_numsig() : Return the number of known signatures.\n"
                            " - get_version() : Return the version of Clamav.\n"
                            " - set_debug() : Enable libclamav debug messages.\n"
                            " - version() : Return the version of pyclamav.\n"));

    /* set up archive limits */
    memset(&limits, 0, sizeof(struct cl_limits));
    limits.maxfiles = 1000; /* max files */
    limits.maxfilesize = 10 * 1048576; /* maximal archived file size == 10 Mb */
    limits.maxreclevel = 5; /* maximal recursion level */
    limits.maxratio = 200; /* maximal compression ratio */
    limits.archivememlim = 0; /* disable memory limit for bzip2 scanner */
}

int main(int argc, char *argv[])
{
    /* Pass argv[0] to the Python interpreter */
    Py_SetProgramName(argv[0]);

    /* Initialize the Python interpreter.  Required. */
    Py_Initialize();

    /* Add a static module */
    initpyclamav();
    return 0;
}
